
package com.wjf.wjfaiagent.advisor;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.model.MessageAggregator;
import reactor.core.publisher.Flux;

import java.nio.file.AccessDeniedException;

/**
 * 自定义权限校验拦截器
 */
@Slf4j
public class PermissionCheckAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {

	private int order;
	@Override
	public String getName() {
		return this.getClass().getSimpleName();
	}

	@Override
	public int getOrder() {
		return this.order;
	}

	private boolean hadPermission(String userRole){
		return "admin".equals(userRole);
	}

	private AdvisedRequest before(AdvisedRequest request) throws AccessDeniedException {
//		log.info("ai request: {}", request.userText());
		String userRole = (String)request.advisorParams().get("userRole");
		log.info("userRole：{}", userRole);
		if (!hadPermission(userRole))
			throw new AccessDeniedException("User has no permission to access AI service");
		return request;
	}

	private void observeAfter(AdvisedResponse advisedResponse) {
		log.info("ai response: {}", advisedResponse.response().getResult().getOutput().getText());
	}

	@Override
	public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {

		try {
			advisedRequest = before(advisedRequest);
		} catch (AccessDeniedException e) {
			throw new RuntimeException(e);
		}

		AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);

		observeAfter(advisedResponse);

		return advisedResponse;
	}

	@Override
	public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {

		try {
			advisedRequest = before(advisedRequest);
		} catch (AccessDeniedException e) {
			throw new RuntimeException(e);
		}

		Flux<AdvisedResponse> advisedResponses = chain.nextAroundStream(advisedRequest);

		return new MessageAggregator().aggregateAdvisedResponse(advisedResponses, this::observeAfter);
	}

}
